import tensorflow as tf
import os

train_data_dir='F:/datasets/tianchi_cv_SVHN/mchar_train/'


def _decode_and_resize(filename, label):
    image_string = tf.io.read_file(filename)  # 读取原始文件
    image_decoded = tf.image.decode_jpeg(image_string)  # 解码JPEG图片
    tf.image.crop_and_resize(image_decoded,)
    image_resized = tf.image.resize(image_decoded, [256, 256]) / 255.0
    return image_resized, label

def prepareTrainData():
    train_filenames = tf.constant([train_data_dir + filename for filename in os.listdir(train_data_dir)])
    train_labels = tf.concat([
        tf.zeros(train_filenames.shape, dtype=tf.int32),
        tf.ones(train_dog_filenames.shape, dtype=tf.int32)],
        axis=-1)

    train_dataset = tf.data.Dataset.from_tensor_slices((train_filenames, train_labels))
    train_dataset = train_dataset.map(
        map_func=_decode_and_resize,
        num_parallel_calls=tf.data.experimental.AUTOTUNE  # 并行生成计算数据
    )
    # 取出前buffer_size个数据放入buffer，并从其中随机采样，采样后的数据用后续数据替换
    train_dataset = train_dataset.shuffle(buffer_size=23000)
    train_dataset = train_dataset.batch(batch_size)
    # prefetch() 方法，使得我们可以让数据集对象 Dataset 在训练时预取出若干个元素，使得在 GPU 训练的同时 CPU 可以准备数据，从而提升训练流程的效率。
    train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)
    # 通过 prefetch() 的使用和在 map() 过程中加入 num_parallel_calls 参数，模型训练的时间可缩减至原来的一半甚至更低。
    return train_dataset

